from class_model import VGG19
import torch
import torchvision.transforms as transforms
from PIL import Image
import cv2
import numpy as np
from math import cos, pi
import matplotlib.pyplot as plt

max_epoch = 2000
def adjust_learning_rate(optimizer, current_epoch,max_epoch,lr_min=0., lr_max=0.01,warmup=True):
    warmup_epoch = 5 if warmup else 0
    lr = optimizer.param_groups[0]['lr']
    lr_min = lr * 0.5
    if current_epoch < warmup_epoch:
        lr = lr_max * current_epoch / warmup_epoch
    else:
        lr = lr_min + 0.5 * (lr_max - lr_min) * (
            1.0
            + cos(
                pi
                * (current_epoch - warmup_epoch)
                / (max_epoch - warmup_epoch)
            )
        )
        # lr = lr *(
        #     0.0
        #     + cos(
        #         pi*0.5
        #         * (current_epoch - warmup_epoch)
        #         / (max_epoch - warmup_epoch)))
    # print(lr)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
 
model = VGG19(num_classes=14, init_weights=False)
optimizer = torch.optim.SGD(model.parameters(), lr=0.0001, momentum=0.9)
x = []
lr = []
for epoch in range(max_epoch):
    adjust_learning_rate(optimizer=optimizer, current_epoch=epoch, max_epoch=max_epoch, lr_min=0.0001, lr_max=0.01, warmup=True)
    x.append(epoch + 1)
    lr.append(optimizer.param_groups[0]['lr'])
plt.scatter(x, lr, s = 5)
plt.show()
plt.savefig("re.png")